// Copyright (C) 2010 and onwards Harik Shazeer Labs
// 
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Noam Shazeer and Georges Harik

#ifndef _JOBQUEUE_H_
#define _JOBQUEUE_H_

#include <pthread.h>
#include <iostream>
#include <vector>
#include "util.h"

using namespace std; 

// The global job queue
struct JobQueue;
extern JobQueue * G_JQ;

// Initializing the global job queue
void InitJobQueue();

void * JobQueueLaunchThread(void *q);

struct JobQueue {
  struct Job;

  struct Task {
    virtual void Run() = 0;
    virtual ~Task() {       
      // if (this < (void *)(1 << 30)) 
      // cout << "destroying " << this << endl;
    }
  Task() :autodelete_(true) {}
    Job *job_;
    bool autodelete_;
  };
  
  struct Job {
    int unfinished_tasks_;
    pthread_cond_t finished_cv_;
    JobQueue *queue_;
    Job(JobQueue *q) { 
      queue_ = q;
      pthread_cond_init(&finished_cv_, NULL);
      unfinished_tasks_ = 0;
    }
    // Create the task outside, but the jobqueue deletes it
    void AddTask(Task *t) {
      pthread_mutex_lock(&queue_->lock_);
      t->job_ = this;
      tasks_.push_back(t);
      unfinished_tasks_++;
      queue_->todo_.push_back(t);
      queue_->task_count_++;
      pthread_cond_signal(&queue_->todo_cv_);
      pthread_mutex_unlock(&queue_->lock_);
    }
    void Wait() {
      queue_->job_count_++;
      queue_->ThreadBody(this);
      CHECK(unfinished_tasks_ == 0);
      /*
      pthread_mutex_lock(&queue_->lock_);
      while (unfinished_tasks_ != 0) {
	// pthread_cond_wait(&finished_cv_, &queue_->lock_);
      }
      pthread_mutex_unlock(&queue_->lock_);
      */
    };
    vector<Task *> tasks_;
  };
    
  JobQueue() {
    pthread_mutex_init(&lock_, NULL);
    pthread_cond_init(&todo_cv_, NULL);
    job_count_ = 0;
    task_count_ = 0;
  }
  void PrintStats() {
    cout << "Job count: " << job_count_ << endl;
    cout << "Task count: " << task_count_ << endl;
  }
  uint64 job_count_;
  uint64 task_count_;
  
  void ThreadBody(Job * wait_for = NULL) {
    while (1) {
      pthread_mutex_lock(&lock_);
      if (wait_for && (wait_for->unfinished_tasks_ == 0)) {
	pthread_mutex_unlock(&lock_);
	return;
      }
      while (todo_.size() == 0) {
	pthread_cond_wait(&todo_cv_, &lock_);
	if (wait_for && (wait_for->unfinished_tasks_ == 0)) {
	  pthread_cond_signal(&todo_cv_);
	  pthread_mutex_unlock(&lock_);
	  return;
	}
      }
      Task *t = todo_.back();
      todo_.pop_back();
      if (todo_.size()) pthread_cond_signal(&todo_cv_);

      pthread_mutex_unlock(&lock_);
      //if (t < (void *)(1 << 30)) 
      // cout << "About to run " << t << endl;
      t->Run();
      pthread_mutex_lock(&lock_);
      t->job_->unfinished_tasks_--;
      if (t->job_->unfinished_tasks_ == 0) {
	// cout << "End of job" << endl;
	pthread_cond_broadcast(&todo_cv_);
	// int sig_ret = pthread_cond_signal(&t->job_->finished_cv_);
	// if (sig_ret != 0) cout << "*** FAILED *** " << sig_ret << endl;
      }
      bool to_delete = t->autodelete_;
      pthread_mutex_unlock(&lock_);
      if (to_delete) {
	delete t;
      }
    }
  };
  
  pthread_mutex_t lock_;
  vector<Task *> todo_;
  pthread_cond_t todo_cv_;
  vector<pthread_t *> threads_;

  int next_thread_num_; // for debugging

  void LaunchThread() {
    pthread_mutex_lock(&lock_);
    //cout << "Launching thread " << endl;
    pthread_t * thread = new pthread_t;
    int rc = pthread_create(thread, NULL, JobQueueLaunchThread, (void *)this);
    //cout << "Launched thread return code " << rc  << endl;
    if (rc) {
      cout << "Error in pthread_create - return code " << rc << endl;
    }
    threads_.push_back(thread);
    pthread_mutex_unlock(&lock_);
  }

  void LaunchThreads(int num_threads) {
    for (int i=0; i<num_threads; i++) LaunchThread();
  }
  
  uint64 NumThreads() { return threads_.size(); }
};

// some common tasks that one might want to launch on the jobqueue

void JobQueueZero(JobQueue * jq, void * location, uint64 bytes);
void JobQueueCopy(JobQueue * jq, uint64 * destination, uint64 *source, 
		  uint64 words);
void JobQueueMove(JobQueue * jq, uint64 * destination, uint64 *source, 
		  uint64 words, uint64 *buffer, uint64 buffer_size);

// copy second part of pairs into the given destination
template<class F, class S> struct CopySecondTask : public JobQueue::Task{
  CopySecondTask(S *d, pair<F, S> *b, pair<F, S> *e) {
    destination_ = d; 
    begin_ = b;
    end_ = e;
  }
  void Run() {
    for (pair<F, S> *run = begin_; run!=end_; run++) {
      *destination_ = run->second;
      destination_++;
    }
  }
  S * destination_;
  pair<F, S> * begin_;
  pair<F, S> * end_;
};


struct CopyTask : public JobQueue::Task {
  CopyTask(uint64 *destination, uint64 *source, uint64 size) {
    destination_ = destination;
    source_ = source;
    size_ = size;
  }
  void Run() {
    memcpy(destination_, source_, size_ * sizeof(uint64));
  }
  uint64 *destination_;
  uint64 *source_;
  uint64 size_;
};

struct TouchTask : public JobQueue::Task {
  void Run() {
    // cout << "Touching " << (end_ - begin_)  << endl;
    sum_ = 0;
    for (uint64 *p=begin_; p<end_; p += 512) {
      sum_ += (uint64)(*p);
    }
  }
 TouchTask(uint64 *begin, uint64 *end) :begin_(begin), end_(end) {};
  uint64 * begin_;
  uint64 * end_;
  uint64 sum_;
};

// to force numa allocation for maximum bandwidth, run this 
// immediately after allocating new memory.  
// alternatively, run numactl --interleave=all <process>
inline void DistributedTouch(void * a, uint64 size) {
  size /= sizeof(uint64);
  JobQueue::Job j(G_JQ);
  int tasks = 1024;
  for (int i=0; i<tasks; i++) {
    uint64 * p = (uint64 *)a;
    j.AddTask(new TouchTask(p + ShardBegin(size, tasks, i), 
			    p + ShardEnd(size, tasks, i)));
  }
  j.Wait();
}

struct CheckSumTask : public JobQueue::Task {
  void Run() {
    uint64 sum = 0;
    char *end = m_.end();
    for (char *p = m_.begin(); p<end; p++)
      sum += (uint8)(*p);
    __sync_fetch_and_add(master_sum_, sum);
  }
 CheckSumTask(Memblock m, uint64* master_sum) 
   :m_(m), master_sum_(master_sum) {}
  Memblock m_;
  uint64 *master_sum_;
};

inline uint64 JQCheckSum(Memblock m) {
  JobQueue::Job j(G_JQ);
  uint64 master_sum = 0;
  int tasks = 1024;
  for (int i=0; i<tasks; i++) {
    CheckSumTask *task =
      new CheckSumTask(Memblock(m.begin() + ShardBegin(m.size_, tasks, i), 
				ShardSize(m.size_, tasks, i)), &master_sum);
    j.AddTask(task);
  }
  j.Wait();
  return master_sum;
}

void TimeTest();


#endif
